import gc
import shutil
from typing import Optional, Any

from fastapi import FastAPI, Body, HTTPException, Request
from fastapi.responses import JSONResponse
from loguru import logger
from tools.utils import set_all_random_seed, get_timbre_file, postprocess, get_demo_audio_path, verify_token, decode_verify_token, download_timbre_file, postprocess_timbre_file
from config import (
    prompt_sr, 
    MODEL_PATH, 
    demo_audio_dir, 
    DEMO_AUDIO_TEXT, 
    USER_VOICE_FOLDER, 
    RANDOM_SEED, 
    DEMO_AUDIO_OSS_DIR,
    GO_SERVICE_SECRET_KEY,
    ALGORITHM
    )
from cosyvoice.utils.file_utils import load_wav
from cosyvoice.cli.cosyvoice import CosyVoice, CosyVoice2
from tools.watermark import AudioWatermark
import torch
import os
import numpy as np
import base64
from io import BytesIO
import soundfile as sf
from sse_starlette.sse import EventSourceResponse
from tools._oss import oss_client
from task_queue2 import TaskManager
from server.db.repository.voice_clone_task_repository import update_voice_clone_task_status, get_voice_clone_task_by_timbre_id
from server.db.base import Base, engine
import redis
from config import REDIS_HOST, REDIS_DB, global_redis_client
import time
import uvicorn

# 使用时传入redis客户端实例
task_manager = TaskManager(max_concurrent_tasks=1, redis_client=global_redis_client)

# 如果数据表不存在，则创建数据表
Base.metadata.create_all(bind=engine)

# 创建一个CosyVoice实例
cosyvoice = CosyVoice2(MODEL_PATH, fp16=False) if "CosyVoice2" in MODEL_PATH else CosyVoice(MODEL_PATH, load_jit=False, load_trt=True, fp16=False)

# 创建一个FastAPI实例
app = FastAPI(title="Generate Audio API", description="Generate Audio API")

# 创建音频水印对象
watermark = AudioWatermark()

# 音频生成器，支持流式返回音频
async def audio_generator(text, prompt_speech_16k, uid, timbre_id, stream, task_id):
    with torch.no_grad():
        full_audio_data = np.array([])
        for i in cosyvoice.inference_cross_lingual(text, prompt_speech_16k, uid, timbre_id, stream=stream):
            audio_data = i['tts_speech'].numpy().flatten()

            full_audio_data = np.concatenate((full_audio_data, audio_data))

            buffer_wav_io = BytesIO()
            sf.write(buffer_wav_io, audio_data, cosyvoice.sample_rate, format="wav")
            buffer_wav_io.seek(0)
            audio_base64 = base64.b64encode(buffer_wav_io.read()).decode('utf-8')

            yield {"event": "audio",
                    "data": {"status_code": 200, "isFinish": False, "audio_data": audio_base64,
                            "task_id": task_id, "sample_rate": cosyvoice.sample_rate}}

        # 给音频添加水印
        _, watermarked_audio = watermark.embed_watermark(full_audio_data, cosyvoice.sample_rate)
        # 检测水印
        detected_watermark = watermark.detect_watermark(watermarked_audio,
                                                        watermark_length=len(watermark.watermark_text))
        logger.debug(f"检测到的水印: {detected_watermark}")
        # 生成完整音频
        buffer_wav_io_full = BytesIO()
        # 保存完整音频
        sf.write(buffer_wav_io_full, watermarked_audio, cosyvoice.sample_rate, format="wav")
        # 读取完整音频
        buffer_wav_io_full.seek(0)
        # 转换为base64编码
        full_audio_base64 = base64.b64encode(buffer_wav_io_full.read()).decode('utf-8')
        # 释放内存
        torch.cuda.empty_cache()
        gc.collect()
        yield {"event": "audio",
                "data": {"status_code": 200, "isFinish": True, "audio_data": full_audio_base64,
                        "task_id": task_id, "sample_rate": cosyvoice.sample_rate}}

def audio_generate_clone(text, prompt_speech_16k, uid, timbre_id, task_id):
    update_voice_clone_task_status(task_id, "processing") # 更新任务状态为处理中
    with torch.no_grad():
        full_audio_data = np.array([])
        for i in cosyvoice.inference_cross_lingual(text, prompt_speech_16k, uid, timbre_id, stream=False):
            audio_data = i['tts_speech'].numpy().flatten()
            full_audio_data = np.concatenate((full_audio_data, audio_data))
        
        # 释放内存
        torch.cuda.empty_cache()
        gc.collect()
        # 给音频添加水印
        _, watermarked_audio = watermark.embed_watermark(full_audio_data, cosyvoice.sample_rate)
        # 检测水印
        detected_watermark = watermark.detect_watermark(watermarked_audio,
                                                        watermark_length=len(watermark.watermark_text))
        logger.debug(f"检测到的水印: {detected_watermark}")
        user_voice_folder = os.path.join(USER_VOICE_FOLDER, uid, timbre_id)  # 用户音色文件夹
        local_demo_audio_dir = os.path.join(demo_audio_dir, uid, timbre_id)  # 本地试听音频目录
        demo_audio_path = get_demo_audio_path(uid=uid, timbre_id=timbre_id, demo_audio_dir=local_demo_audio_dir) # 本地试听音频路径
        try:
            sf.write(demo_audio_path, watermarked_audio, cosyvoice.sample_rate, format="wav")
        except Exception as e:
            logger.error(f"写入音频文件时出错: {str(e)}")
            task_manager.fail_task(task_id)
            raise HTTPException(status_code=502, detail="写入音频文件时发生问题。")
        try:
            oss_client.upload_file(local_file_path=demo_audio_path,
                                            directory=DEMO_AUDIO_OSS_DIR)  # 上传试听音频文件到OSS
            if os.path.exists(local_demo_audio_dir):
                shutil.rmtree(local_demo_audio_dir)
            logger.debug(f"音频文件已上传到OSS: {demo_audio_path}")
            update_voice_clone_task_status(task_id, "completed")  # 更新任务状态为已完成
            task_manager.remove_task(task_id)  # 任务执行完毕，删除活跃队列中的任务
        except Exception as e:
            if os.path.exists(user_voice_folder):
                shutil.rmtree(user_voice_folder)
            if os.path.exists(local_demo_audio_dir):
                shutil.rmtree(local_demo_audio_dir)
            task_manager.fail_task(task_id)
            logger.error(f"上传音频文件时出错: {str(e)}")
            raise HTTPException(status_code=502, detail="上传音频文件时出错。")
            

# API: 生成TTS音频
@app.post("/generate_audio", summary="MarsX AI TTS")
async def tts(uid: str = Body(..., description="用户ID"),
              timbre_id: str = Body(..., description="音色ID"),
              text: Any = Body(None, description="文本内容"),
              task_id: str = Body(..., description="任务ID"),
              stream: bool = Body(False, description="是否流式生成"),
              seed: Optional[int] = Body(RANDOM_SEED, description="随机种子"),
              task_type: Optional[str] = Body("tts", description="任务类型")):
    """
    生成TTS音频。

    :param uid: 用户ID
    :param timbre_id: 音色ID
    :param text: 文本内容
    :param task_id: 任务ID
    :param stream: 是否流式生成
    :param seed: 随机种子
    :return: EventSourceResponse 实时返回生成的音频片段
    """
    # 获取并验证提示音色文件
    try:
        prompt_wav = get_timbre_file(uid, timbre_id)
    except Exception as e:
        logger.error(f"获取音色文件时出现错误: {str(e)}")
        return JSONResponse(status_code=200,
                            content={"status_code": 400, "isSuccess": False, "message": "获取音色文件时出现错误。",
                                     "task_id": task_id})

    if not prompt_wav:
        logger.debug(f"未找到音色文件，从MySQL中获取音色 {timbre_id} 的URL")
        voice_clone_task = get_voice_clone_task_by_timbre_id(timbre_id)
        if not voice_clone_task:
            return JSONResponse(status_code=200,
                                content={"status_code": 400, "isSuccess": False,
                                         "message": f"MySQL中未找到音色ID为 {timbre_id} 的音色数据。",
                                         "task_id": task_id})
        try:
            download_result, local_file_path = await download_timbre_file(uid, timbre_id) # 下载音色文件
            # 如果下载失败，返回错误信息
            if not download_result:
                return JSONResponse(status_code=200,
                                    content={"status_code": 400, "isSuccess": False,
                                            "message": "下载音色文件时出现错误。",
                                            "task_id": task_id})
            # 如果格式转换失败，返回错误信息
            if not postprocess_timbre_file(local_file_path): # 格式转换
                return JSONResponse(status_code=200,
                                    content={"status_code": 400, "isSuccess": False,
                                            "message": "音色文件格式转换时出现错误。",
                                            "task_id": task_id})
            
        except Exception as e:
            logger.error(f"获取音色文件时出现错误: {str(e)}")
            return JSONResponse(status_code=200,
                                content={"status_code": 400, "isSuccess": False,
                                         "message": "获取音色文件时出现错误。",
                                         "task_id": task_id})
        prompt_wav = get_timbre_file(uid, timbre_id)
        if not prompt_wav:
            logger.error(f"再次获取音色文件失败: {timbre_id}")
            return JSONResponse(status_code=200,
                                content={"status_code": 400, "isSuccess": False, "message": "再次获取音色文件失败。",
                                         "task_id": task_id})

    # 处理提示语音
    try:
        prompt_speech_16k = postprocess(load_wav(prompt_wav, prompt_sr), cosyvoice.sample_rate)
    except Exception as e:
        if task_type == "tts":
            logger.error(f"语音合成任务，预处理音色文件时出现错误: {str(e)}")
            return JSONResponse(status_code=200,
                                content={"status_code": 400, "isSuccess": False,
                                         "message": "语音合成任务，预处理音色文件时出错。",
                                         "task_id": task_id})
        if task_type == "clone":
            logger.error(f"声音克隆任务，预处理音色文件时出现错误: {str(e)}")
            task_manager.fail_task(task_id)
            raise HTTPException(status_code=500, detail="预处理音色文件时出现错误。")

    set_all_random_seed(seed)

    if task_type == "tts":

        return EventSourceResponse(audio_generator(text, prompt_speech_16k, uid, timbre_id, stream, task_id))

    if task_type == "clone":
        audio_generate_clone(DEMO_AUDIO_TEXT, prompt_speech_16k, uid, timbre_id, task_id)


# API: 生成TTS音频
@app.post("/generate_audio_token", summary="MarsX AI TTS")
async def tts(uid: str = Body(..., description="用户ID"),
              timbre_id: str = Body(..., description="音色ID"),
              text: Any = Body(None, description="文本内容"),
              task_id: str = Body(..., description="任务ID"),
              stream: bool = Body(False, description="是否流式生成"),
              seed: Optional[int] = Body(RANDOM_SEED, description="随机种子"),
              task_type: Optional[str] = Body("tts", description="任务类型"),
              request: Request = None):
    """
    生成TTS音频。

    :param uid: 用户ID
    :param timbre_id: 音色ID
    :param text: 文本内容
    :param task_id: 任务ID
    :param stream: 是否流式生成
    :param seed: 随机种子
    :return: EventSourceResponse 实时返回生成的音频片段
    """
    # 验证Token
    verify_token_result = verify_token(request.headers.get("Authorization"), "sub", timbre_id)
    if verify_token_result["status"] != "0":
        return JSONResponse(status_code=200,
                            content={"status_code": 400, "isSuccess": False, "message": verify_token_result["message"],
                                     "task_id": task_id})
    # 获取并验证提示音色文件
    try:
        prompt_wav = get_timbre_file(uid, timbre_id)
    except Exception as e:
        logger.error(f"获取音色文件时出现错误: {str(e)}")
        return JSONResponse(status_code=200,
                            content={"status_code": 400, "isSuccess": False, "message": "获取音色文件时出现错误。",
                                     "task_id": task_id})

    if not prompt_wav:
        logger.debug(f"未找到音色文件，从MySQL中获取音色 {timbre_id} 的URL")
        voice_clone_task = get_voice_clone_task_by_timbre_id(timbre_id)
        if not voice_clone_task:
            return JSONResponse(status_code=200,
                                content={"status_code": 400, "isSuccess": False,
                                         "message": f"MySQL中未找到音色ID为 {timbre_id} 的音色数据。",
                                         "task_id": task_id})
        try:
            download_result, local_file_path = await download_timbre_file(uid, timbre_id) # 下载音色文件
            if not download_result:
                return JSONResponse(status_code=200,
                                    content={"status_code": 400, "isSuccess": False,
                                             "message": "下载音色文件时出现错误。",
                                             "task_id": task_id})
            if not postprocess_timbre_file(local_file_path): # 格式转换
                return JSONResponse(status_code=200,
                                    content={"status_code": 400, "isSuccess": False,
                                             "message": "音色文件格式转换时出现错误。",
                                             "task_id": task_id})
        except Exception as e:
            logger.error(f"获取音色文件时出现错误: {str(e)}")
            return JSONResponse(status_code=200,
                                content={"status_code": 400, "isSuccess": False,
                                         "message": "获取音色文件时出现错误。",
                                         "task_id": task_id})
        prompt_wav = get_timbre_file(uid, timbre_id)
        if not prompt_wav:
            logger.error(f"再次获取音色文件失败: {timbre_id}")
            return JSONResponse(status_code=200,
                                content={"status_code": 400, "isSuccess": False, "message": "再次获取音色文件失败。",
                                         "task_id": task_id})

    # 处理提示语音
    try:
        prompt_speech_16k = postprocess(load_wav(prompt_wav, prompt_sr), cosyvoice.sample_rate)
    except Exception as e:
        if task_type == "tts":
            logger.error(f"语音合成任务，预处理音色文件时出现错误: {str(e)}")
            return JSONResponse(status_code=200,
                                content={"status_code": 400, "isSuccess": False,
                                         "message": "语音合成任务，预处理音色文件时出错。",
                                         "task_id": task_id})
        if task_type == "clone":
            logger.error(f"声音克隆任务，预处理音色文件时出现错误: {str(e)}")
            task_manager.fail_task(task_id)
            update_voice_clone_task_status(task_id, "failed")
            raise HTTPException(status_code=500, detail="预处理音色文件时出现错误。")

    set_all_random_seed(seed)

    if task_type == "tts":
        

        return EventSourceResponse(audio_generator(text, prompt_speech_16k, uid, timbre_id, stream, task_id))

    if task_type == "clone":
        audio_generate_clone(DEMO_AUDIO_TEXT, prompt_speech_16k, uid, timbre_id, task_id)


@app.post("/v2/generate_audio", summary="MarsX AI TTS")
async def tts(uid: str = Body(..., description="用户ID"),
              timbre_id: str = Body(..., description="音色ID"),
              isDefault: bool = Body(True, description="是否使用默认音色"),
              text: Any = Body(None, description="文本内容"),
              task_id: str = Body(..., description="任务ID"),
              stream: bool = Body(False, description="是否流式生成"),
              seed: Optional[int] = Body(RANDOM_SEED, description="随机种子"),
              task_type: Optional[str] = Body("tts", description="任务类型"),
              request: Request = None):
    """
    生成TTS音频。

    :param uid: 用户ID
    :param timbre_id: 音色ID
    :param text: 文本内容
    :param isDefault: 是否使用默认音色
    :param task_id: 任务ID
    :param stream: 是否流式生成
    :param seed: 随机种子
    :return: EventSourceResponse 实时返回生成的音频片段
    """
    # 验证Token
    try:
        token = decode_verify_token(request, GO_SERVICE_SECRET_KEY, ALGORITHM)
    except Exception as e:
        return JSONResponse(status_code=200,
                            content={"status_code": 401, "isSuccess": False, "message": f"解析Token时出现错误: {e.detail.get('msg')}",
                                     "task_id": task_id})
    # 判断是否使用默认音色
    if isDefault:
        uid = "pre"  # 如果使用默认音色，则将用户ID设置为"pre"，因为默认音色的UID为"pre"
    # 获取并验证提示音色文件
    try:
        prompt_wav = get_timbre_file(uid, timbre_id)
    except Exception as e:
        logger.error(f"获取音色文件时出现错误: {str(e)}")
        return JSONResponse(status_code=200,
                            content={"status_code": 400, "isSuccess": False, "message": "获取音色文件时出现错误。",
                                     "task_id": task_id})

    if not prompt_wav:
        logger.debug(f"未找到音色文件，从MySQL中获取音色 {timbre_id} 的URL")
        voice_clone_task = get_voice_clone_task_by_timbre_id(timbre_id)
        if not voice_clone_task:
            return JSONResponse(status_code=200,
                                content={"status_code": 400, "isSuccess": False,
                                         "message": f"MySQL中未找到音色ID为 {timbre_id} 的音色数据。",
                                         "task_id": task_id})
        try:
            download_result, local_file_path = await download_timbre_file(uid, timbre_id) # 下载音色文件
            if not download_result:
                return JSONResponse(status_code=200,
                                    content={"status_code": 400, "isSuccess": False,
                                             "message": "下载音色文件时出现错误。",
                                             "task_id": task_id})
            if not postprocess_timbre_file(local_file_path): # 格式转换
                return JSONResponse(status_code=200,
                                    content={"status_code": 400, "isSuccess": False,
                                             "message": "音色文件格式转换时出现错误。",
                                             "task_id": task_id})
        except Exception as e:
            logger.error(f"获取音色文件时出现错误: {str(e)}")
            return JSONResponse(status_code=200,
                                content={"status_code": 400, "isSuccess": False,
                                         "message": "获取音色文件时出现错误。",
                                         "task_id": task_id})
        prompt_wav = get_timbre_file(uid, timbre_id)
        if not prompt_wav:
            logger.error(f"再次获取音色文件失败: {timbre_id}")
            return JSONResponse(status_code=200,
                                content={"status_code": 400, "isSuccess": False, "message": "再次获取音色文件失败。",
                                         "task_id": task_id})

    # 处理提示语音
    try:
        prompt_speech_16k = postprocess(load_wav(prompt_wav, prompt_sr), cosyvoice.sample_rate)
    except Exception as e:
        if task_type == "tts":
            logger.error(f"语音合成任务，预处理音色文件时出现错误: {str(e)}")
            return JSONResponse(status_code=200,
                                content={"status_code": 400, "isSuccess": False,
                                         "message": "语音合成任务，预处理音色文件时出错。",
                                         "task_id": task_id})
        if task_type == "clone":
            logger.error(f"声音克隆任务，预处理音色文件时出现错误: {str(e)}")
            task_manager.fail_task(task_id)
            update_voice_clone_task_status(task_id, "failed")
            raise HTTPException(status_code=500, detail="预处理音色文件时出现错误。")

    set_all_random_seed(seed)

    if task_type == "tts":
        return EventSourceResponse(audio_generator(text, prompt_speech_16k, uid, timbre_id, stream, task_id))

    if task_type == "clone":
        audio_generate_clone(DEMO_AUDIO_TEXT, prompt_speech_16k, uid, timbre_id, task_id)
